//----------------------------------*-C++-*----------------------------------//
/*!
 * \file   Newton.cc
 * \author Jeremy Roberts
 * \date   11/24/2010
 * \brief  Member definitions of base class Newton
 * \note   Copyright (C) 2010 Jeremy Roberts
 */
//---------------------------------------------------------------------------//
// $Rev::                                               $:Rev of last commit
// $Author::                                            $:Author of last commit
// $Date::                                              $:Date of last commit
//---------------------------------------------------------------------------//

#include <iostream>
#include <cmath>
#include "petscsnes.h"
#include "LinAlg.hh"
#include "GlobalInput.hh"
#include "ResponseFunctionServer.hh"
#include "ResponseMatrix.hh"
#include "ResponseMatrixFull.hh"
#include "AbsorptionResponse.hh"
#include "FissionResponse.hh"
#include "LeakageResponse.hh"
#include "ConnectMatrix.hh"
#include "Connect2dCart.hh"
#include "GlobalProblem.hh"
#include "GlobalSolver.hh"
#include "Newton.hh"
#include "ResidualWrap.hh"
#include "JacobianShell.hh"
#include "JacobianEmpty.hh"
#include "PCAppxJacobian.hh"
//#include "PIPCShell.hh"


//---------------------------------------------------------------------------//
// CONSTRUCTOR
//---------------------------------------------------------------------------//


Newton::Newton( GlobalProblem &pr, GlobalInput &in )
 : GlobalSolver(pr,in)
{
    std::cout << " CONSTRUCTING Newton " << std::endl;
    // nothing more here for now
}


//---------------------------------------------------------------------------//
// DESTRUCTOR
//---------------------------------------------------------------------------//
Newton::~Newton()
{
    // nothing more here right now, but should write Newton->destroy();
    return; 
}

//---------------------------------------------------------------------------//
/*!
 * \brief This is the Newton solver.
 *
 */
void Newton::Solve()
{

    scalar one = 1.0, zero = 0.0;

    // - - - - - - - - - -- - - - - - - - - - - - - - - - - - - - - - - - - - 
    //  Create nonlinear solver context
    // - - - - - - - - - -- - - - - - - - - - - - - - - - - - - - - - - - - - 
    SNESCreate( PETSC_COMM_WORLD, &snes );

    // - - - - - - - - - -- - - - - - - - - - - - - - - - - - - - - - - - - - 
    //  Create vectors for solution and nonlinear function
    // - - - - - - - - - -- - - - - - - - - - - - - - - - - - - - - - - - - - 
    x = new SermentVector( input->degfree+2 );
    f = new SermentVector( input->degfree+2 );
    // initialize
    for ( integer i=0; i<input->degfree+2; i++ ){
        x->insertVal( i, zero );
        f->insertVal( i, one );
    }
    // uniform zeroth order initial guess
    for ( integer i=0; i<input->degfree; 
          i = i + (input->spaceord+1)*(input->angleord+1)*input->numgroups )
    {
        x->insertVal( i, one );
    }
    x->vecScale(1.0/sqrt(x->vecDot(*x)));        // norm'd uniform 0-order guess
    x->insertVal( input->degfree, input->keff ); // keff guess from input
    x->insertVal( input->degfree+1, one );       // guess for lambda is unity

    // perform one crude power iteration as initial seed
    scalar keff   = input->keff;  
    scalar lambda = 0; 
    scalar *x_a; 
    VecGetArray( x->V, &x_a ); 
    SermentVector  Jinc( input->degfree ); 
    SermentVector  temp( input->degfree );
    VecPlaceArray( Jinc.V, x_a );
    for (int powit = 0; powit < 1; powit++)
    {
        problem->R.updateData( keff ); 
        problem->L.updateData( keff );
        problem->A.updateData( keff ); 
        problem->F.updateData( keff ); 
        for (integer i=0; i<1000; i++)
        {
            problem->R.matVec(Jinc,temp); problem->M.matVec(temp,Jinc);
            lambda = sqrt( Jinc.vecDot(Jinc) ); Jinc.vecScale(1.0/lambda);
        }
        keff = problem->F.vecDot(Jinc)  / 
         (problem->A.vecDot(Jinc) + problem->L.computeLeakage(Jinc) );
    }
    VecResetArray( Jinc.V ); 
    x_a[input->degfree] = keff; 
    x_a[input->degfree+1] = lambda;    
    VecRestoreArray( x->V, &x_a );
    Jinc.releaseMe();
    temp.releaseMe();

    // - - - - - - - - - -- - - - - - - - - - - - - - - - - - - - - - - - - - 
    //  Set the residual and vector in SNES
    // - - - - - - - - - -- - - - - - - - - - - - - - - - - - - - - - - - - - 
    SNESSetFunction( snes, f->V, ResidualWrap, this );

    // - - - - - - - - - -- - - - - - - - - - - - - - - - - - - - - - - - - - 
    //  Set the solver context and solve the problem
    // - - - - - - - - - -- - - - - - - - - - - - - - - - - - - - - - - - - - 
    KSP ksp;
    PC  pc;
    SNESGetKSP(snes,&ksp);
    KSPSetType(ksp,KSPGMRES);

    // - - - - - - - - - -- - - - - - - - - - - - - - - - - - - - - - - - - - 
    //  Set the preconditioner
    // - - - - - - - - - -- - - - - - - - - - - - - - - - - - - - - - - - - - 

    KSPGetPC( ksp, &pc );

    PCAppxJacobian *JacPC;
//    PIPCShell *PIPC;


    if ( input->pctype == 0 ) // no preconditioner
    {
        PCSetType(pc,PCNONE);
    }
    else if ( input->pctype == 1 ) // build the approximate jacobian matrix
    {
        // the approximate jacobian is sparse except for the final two rows
        // and columns.  hence, it makes sense to allocate using the nzz
        // rather than nz parameter.  For the main block, each row has at most
        // degfree/numel + 2 nonzeros;  For the m-1th, m-1, and mth, m-2.
        integer nzz[ input->degfree+2 ];
        for (integer i=0; i < input->degfree; i++)
            nzz[i] = input->degfree/input->numel + 2;
        nzz[input->degfree] = input->degfree+1;
        nzz[input->degfree+1] = input->degfree;
        
//        for (integer i=0; i < input->degfree+2; i++)
//            cout << " NZZ = " << nzz[i] << endl;
        const integer *nzzp = nzz;
        JacPC = new PCAppxJacobian( input->degfree+2, input->degfree+2, 3, 
                           nzzp, x, problem );
        SNESSetLagPreconditioner(snes,-1); // -1 indicates we never rebuild
        PCSetType(pc,PCILU);
        PCFactorSetLevels(pc,input->ilulevel);
    }
    else if (  input->pctype == 2 ) // PI preconditioner
    {

//        cout << " pc type 2 " << endl;
//        PIPC = new PIPCShell(input->degfree+2,
//                             input->degfree+2,PETSC_NULL,x,f,problem,snes);
//        PCSetType(pc,PCSHELL);
//        PCShellSetContext(pc,PIPC);
//        PCShellSetApply(pc,PIPCMatVecWrap);
//   //     KSPSetPreconditionerSide(ksp,PC_RIGHT);
    }


    // - - - - - - - - - -- - - - - - - - - - - - - - - - - - - - - - - - - - 
    //  Create Jacobian matrix data structure
    // - - - - - - - - - -- - - - - - - - - - - - - - - - - - - - - - - - - - 
    Mat Jac;
    integer jj = 0;
    if ( jj == 0 ) 
    {
        fp = new JacobianShell(input->degfree+2,
                               input->degfree+2,this,x,f,problem,snes);
        if ( input->pctype  == 0 )
            SNESSetJacobian( snes, fp->M, fp->M, JacobianEmpty, this );
        else if ( input->pctype == 1 )
            SNESSetJacobian( snes, fp->M, JacPC->M, JacobianEmpty, this );
     //   else if ( input->pctype == 2 )
         //   SNESSetJacobian( snes, fp->M, PIPC->M, JacobianEmpty, this );
    }
    else if ( jj == 1 )
    {
        // testing with matrix free finite difference
        MatCreateSNESMF(snes,&Jac);
        MatMFFDSetFunction(Jac,
                (PetscErrorCode (*)(void*, Vec, Vec))SNESComputeFunction,snes);
        SNESSetJacobian(snes,Jac,JacPC->M,MatMFFDComputeJacobian,this);
    }
    else
    {
        // testing with full finite difference jacobian
        MatCreate( PETSC_COMM_SELF, &Jac );
        MatSetSizes( Jac, PETSC_DECIDE, PETSC_DECIDE, 
                     input->degfree+2, input->degfree+2 );
        MatSetFromOptions( Jac );
        SNESSetJacobian( snes, Jac, Jac, SNESDefaultComputeJacobian, this );
    }

    // use the following to quit after just one A*v.  Usually quits newton, too
    //KSPSetTolerances( ksp, 1.e-5, PETSC_DEFAULT, PETSC_DEFAULT, 1 );

    integer its, itslin;
    scalar tol = 1e-8;
    SNESSetTolerances(snes,tol,tol,tol,10,1000);
    SNESSetFromOptions(snes);
    SNESLineSearchSet(snes,SNESLineSearchNo,PETSC_NULL);
    SNESSolve(snes,PETSC_NULL,x->V);

    SNESGetIterationNumber(snes,&its);
    SNESGetLinearSolveIterations(snes,&itslin);


	// test fissionrates
     SermentVector Jo( x->Length()-2 );
	scalar *j_a;
	VecGetArray( x->V, &x_a );  // get the array	
	VecGetArray( Jo.V, &j_a );  // get the array	
	for ( int i = 0; i < Jo.Length(); i ++ )
		j_a[i] = x_a[i];
	VecRestoreArray( x->V, &x_a );  // get the array	
	VecRestoreArray( Jo.V, &j_a );  // get the array		
    fissionRates(&Jo);

    cout << " FINAL NEWTON EIGENVALUES: " << endl;
    printf (" **** FINAL KEFF   = %12.9f \n", x_a[input->degfree] );
    printf (" **** FINAL LAMBDA = %12.9f \n", x_a[input->degfree+1] );
    printf (" **** NONLINEAR ITERATIONS  = %8i \n", its );
    printf (" **** LINEAR ITERATIONS     = %8i \n", itslin );

    if ( jj == 2 ) // debug, prints fd-jacobian to ascii and binary
    {
        PetscViewer viewer;
        PetscViewerASCIIOpen(PETSC_COMM_WORLD,"jac.output",&viewer);
        PetscViewerSetFormat(viewer, PETSC_VIEWER_ASCII_DENSE );
        MatView(Jac,viewer);
        MatView( Jac, PETSC_VIEWER_STDOUT_SELF );
        PetscViewer bview;
        PetscViewerBinaryOpen(PETSC_COMM_WORLD,"jac.bin",FILE_MODE_WRITE,&bview);
        MatView(Jac,bview);
    }

    return;
}

//---------------------------------------------------------------------------//
/*!
 * \brief This function computes the nonlinear residual.
 *
 * The nonlinear residual is defined as
 *   \f[
 *       \mathbf{f(x)} = \left [\begin{array}{c}
 *	        (\mathbf{M}\mathbf{R}(k)-\lambda \mathbf{I}) \mathbf{J_-} \\
 *	        \mathbf{F}(k)\mathbf{J_-} - (k\mathbf{L}(k)\mathbf{J_-} ) \\
 *	        \frac{1}{2} \mathbf{J^T_-} \mathbf{J_-} - \frac{1}{2}  
 *	      \end{array} 
 *       \right ]  = \mathbf{0} \, ,
 *   \f]
 * which is the same as used in the
 *
 */
PetscErrorCode Newton::Residual( SNES snes, Vec X, Vec F, void *ptr )
{

    scalar *x_a, *f_a;
    VecGetArray( X, &x_a );
    VecGetArray( F, &f_a );

    SermentVector  Jinc( input->degfree );
    SermentVector  F_a( input->degfree );
    SermentVector  temp( input->degfree );

    scalar k, lambda;

    // trying vecplacearray with x_a and f_a into smaller arrays
    VecPlaceArray( Jinc.V, x_a );
    VecPlaceArray( F_a.V, f_a );

    k      = x_a[input->degfree];
    lambda = x_a[input->degfree+1];

    problem->R.updateData( k );
    problem->L.updateData( k );
    problem->F.updateData( k );
    problem->A.updateData( k );

    // first m-2 entries
    problem->R.matVec(Jinc,temp);
    problem->M.matVec(temp,F_a);
    F_a.vecAYPV( -lambda, Jinc ); 

    // (m-1)th entry
    f_a[input->degfree] = problem->F.vecDot( Jinc ) -
                            k*( problem->A.vecDot( Jinc ) +
                                problem->L.computeLeakage( Jinc ) );
    // (m)th entry
    f_a[input->degfree+1] = 0.5 - 0.5*Jinc.vecDot(Jinc);

    VecResetArray( Jinc.V );
    VecResetArray( F_a.V );

    VecRestoreArray( X, &x_a );
    VecRestoreArray( F, &f_a   );

    Jinc.releaseMe();
    F_a.releaseMe();
    temp.releaseMe();

    return 0;
}

//---------------------------------------------------------------------------//
//                 end of Newton.cc
//---------------------------------------------------------------------------//

